﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Cloo;
using System.Diagnostics.Contracts;
using System.Diagnostics;

namespace SimpleCLML.LogisticRegression
{
    public sealed class MulticlassLogisticRegressionTrainer
    {
        private List<MulticlassTrainingExample> examples;
        private int classCount;
        private int featureCount;
        private int exampleCount;

        public MulticlassLogisticRegressionTrainer(IEnumerable<MulticlassTrainingExample> inputExamples, int classCount)
        {
            this.classCount = classCount;
            this.examples = inputExamples.ToList();
            this.featureCount = examples.First().x.Length + 1;
            this.exampleCount = examples.Count;
        }

        public MulticlassLogisticRegression Train(int maxIterations=1280, float errorThreshold=1.0e-6f, float learningRate=0.5f)
        {
            const int ExampleGroupingConstant = 256; // TODO: Get from algorithm
            int alignedFeatureCount = featureCount % 4 == 0 ? featureCount : featureCount + 4 - featureCount % 4;
            int alignedExampleCount = exampleCount % ExampleGroupingConstant == 0 ? exampleCount : exampleCount + ExampleGroupingConstant - exampleCount % ExampleGroupingConstant;

            /* matrices in column-major order */
            float[] X = new float[alignedFeatureCount * alignedExampleCount];
            float[] Y = new float[alignedExampleCount];
            float[] theta0 = new float[alignedFeatureCount];
            float[][] thetas = new float[exampleCount][];

            for (int i = 0; i < examples.Count; ++i)
            {
                var example = examples[i];
                X[i] = 1.0f; // bias
                for (int j = 1; j < featureCount; j++)
                {
                    X[i + j * alignedExampleCount] = example.x[j - 1];
                }
            }

            for (int c = 0; c < classCount; c++)
            {
                Trace.WriteLine(String.Format("Training class {0} of {1}", c + 1, classCount), "Information");
                for (int i = 0; i < exampleCount; ++i)
                {
                    var example = examples[i % examples.Count];
                    Y[i] = example.y == c ? 1.0f : 0.0f;
                }

                // TODO: run multiclass in a single session, without reloading X
                using (var gd = new LogisticRegressionGradientDescent(maxIterations, errorThreshold, learningRate))
                {
                    gd.Init(X, Y, exampleCount, null);
                    float[] theta = gd.Run();
                    Array.Resize<float>(ref theta, featureCount); // cut extra items
                    thetas[c] = theta;
                }
            }

            return new MulticlassLogisticRegression(classCount, thetas);
        }

    }
}
